import numpy as np
import tensorflow as tf
bfloat16 = tf.bfloat16.as_numpy_dtype
dtype_emu = {bfloat16: 0, np.float16: 1, np.float32: 2, np.int8: 3, np.int16: 4, np.int32: 5}

def gen_golden_data_simple():
    
    mask_type = np.int8
    src_type = bfloat16
    # src_type = np.float16
    # src_type = np.float32
    # src_type = np.int8
    # src_type = np.int16
    # src_type = np.int32

    ## 核间均分，单核计算量对齐:
    # input_shape = [32, 1024]

    ## 核间均分，单核计算量非对齐:
    # input_shape = [32, 1023]

    ## 核间不均分，单核计算量对齐:
    # input_shape = [8, 1024]

    ## 核间不均分，单核计算量非对齐:
    input_shape = [8, 1023]
    x = np.random.uniform(-10, 10, input_shape).astype(src_type)
    y = np.random.uniform(-10, 10, input_shape).astype(src_type)
    b = np.random.randint(0, 2, input_shape).astype(mask_type)
    golden = np.where(b, x, y).astype(src_type)
    tiling = np.array([input_shape[0] * input_shape[1], dtype_emu[src_type], dtype_emu[mask_type]], dtype=np.uint32)
    tiling.tofile("./input/input_tiling.bin")
    x.astype(src_type).tofile("./input/input_x.bin")
    y.astype(src_type).tofile("./input/input_y.bin")
    b.astype(mask_type).tofile("./input/input_b.bin")
    golden.astype(src_type).tofile("./output/golden.bin")


if __name__ == "__main__":
    gen_golden_data_simple()
